#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pymupdf_tools.py 模块的测试文件
"""

import os

import fitz
import numpy as np

from app.utils.pymupdf_tools import (
    get_pages_from_range,
    pdf_get_content_region,
    convert_to_images,
    page_to_img,
    pdf_pages_to_images
)
from tests.base_test import base_test_case

logger = base_test_case.get_logger(__name__)
TEST_DATA_DIR = base_test_case.test_data_dir
OUTPUT_DATA_DIR = base_test_case.output_data_dir

# 使用指定的测试PDF文件
TEST_PDF_PATH = TEST_DATA_DIR / "1715339805571.pdf"


def get_pages_from_range_test():
    """测试 get_pages_from_range 函数"""
    # 检查测试PDF文件是否存在
    if not os.path.exists(TEST_PDF_PATH):
        logger.warning(f"测试文件 {TEST_PDF_PATH} 不存在，跳过测试")
        return

    # 打开文档进行测试
    doc = fitz.open(str(TEST_PDF_PATH))
    try:
        # 测试从 Document 对象加载所有页面
        pages = get_pages_from_range(doc)
        assert len(pages) > 0

        # 测试使用范围元组
        pages = get_pages_from_range(doc, (0, min(2, len(pages) - 1)))
        assert len(pages) > 0

        # 测试使用索引列表
        pages = get_pages_from_range(doc, [0, min(1, len(pages) - 1)])
        assert len(pages) > 0
    finally:
        doc.close()

    # 测试从文件路径加载所有页面
    pages = get_pages_from_range(str(TEST_PDF_PATH))
    assert len(pages) > 0


def convert_to_images_test():
    """测试 convert_to_images 函数"""
    # 检查测试PDF文件是否存在
    if not os.path.exists(TEST_PDF_PATH):
        logger.warning(f"测试文件 {TEST_PDF_PATH} 不存在，跳过测试")
        return

    # 测试从 PDF 文件路径转换
    # 注意：当直接传入文件路径时，convert_to_images 内部会处理文档的打开和关闭
    images, failed = convert_to_images(str(TEST_PDF_PATH), zoom_factor=0.5)
    assert len(images) > 0
    assert isinstance(images[0], np.ndarray)
    assert len(failed) == 0

    # 测试从页面对象转换（需要确保文档在使用期间保持打开状态）
    doc = fitz.open(str(TEST_PDF_PATH))
    try:
        pages = get_pages_from_range(doc, [0])  # 只获取第一页
        images, failed = convert_to_images(pages, zoom_factor=0.5)
        assert len(images) > 0
        assert isinstance(images[0], np.ndarray)
    finally:
        doc.close()


def page_to_img_test():
    """测试 page_to_img 函数"""
    # 检查测试PDF文件是否存在
    if not os.path.exists(TEST_PDF_PATH):
        logger.warning(f"测试文件 {TEST_PDF_PATH} 不存在，跳过测试")
        return

    # 打开并测试页面转图像（需要确保文档在使用期间保持打开状态）
    doc = fitz.open(str(TEST_PDF_PATH))
    try:
        page = doc[0]

        # 测试 RGB 模式
        img_rgb = page_to_img(page, zoom_factor=0.5, color_mode="RGB")
        assert isinstance(img_rgb, np.ndarray)
        assert len(img_rgb.shape) == 3  # 应该是3通道

        # 测试 BGR 模式
        img_bgr = page_to_img(page, zoom_factor=0.5, color_mode="BGR")
        assert isinstance(img_bgr, np.ndarray)
        assert len(img_bgr.shape) == 3  # 应该是3通道
    finally:
        doc.close()


def pdf_pages_to_images_test():
    """测试 pdf_pages_to_images 函数"""
    # 检查测试PDF文件是否存在
    if not os.path.exists(TEST_PDF_PATH):
        logger.warning(f"测试文件 {TEST_PDF_PATH} 不存在，跳过测试")
        return

    # 加载页面（需要确保文档在使用期间保持打开状态）
    doc = fitz.open(str(TEST_PDF_PATH))
    try:
        pages = get_pages_from_range(doc, [0])  # 只测试第一页

        # 转换为图像
        images = pdf_pages_to_images(pages, zoom_factor=0.5)
        assert len(images) > 0
        assert isinstance(images[0], np.ndarray)
    finally:
        doc.close()


def pdf_get_content_region_test():
    """测试 pdf_get_content_region 函数"""
    # 检查测试PDF文件是否存在
    if not os.path.exists(TEST_PDF_PATH):
        logger.warning(f"测试文件 {TEST_PDF_PATH} 不存在，跳过测试")
        return

    # 测试从页面对象检测内容区域（需要确保文档在使用期间保持打开状态）
    doc = fitz.open(str(TEST_PDF_PATH))
    try:
        page = doc[0]

        region = pdf_get_content_region(page, zoom_factor=1.0)
        # 检查返回值是否为有效的区域坐标元组
        if region is not None:
            assert isinstance(region, tuple)
            assert len(region) == 4
            x0, y0, x1, y1 = region
            assert x0 >= 0 and y0 >= 0 and x1 > x0 and y1 > y0
    finally:
        doc.close()


if __name__ == '__main__':
    # 检查测试PDF文件是否存在
    if not os.path.exists(TEST_PDF_PATH):
        print(f"警告: 测试文件 {TEST_PDF_PATH} 不存在，请确认测试数据文件已准备就绪")
    else:
        # 运行所有测试
        get_pages_from_range_test()
        print("✓ get_pages_from_range_test passed")

        convert_to_images_test()
        print("✓ convert_to_images_test passed")

        page_to_img_test()
        print("✓ page_to_img_test passed")

        pdf_pages_to_images_test()
        print("✓ pdf_pages_to_images_test passed")

        pdf_get_content_region_test()
        print("✓ pdf_get_content_region_test passed")

        print("All tests passed!")
